
package net.pragyah.scalby.inference.exact

import net.pragyah.scalby.Variable



import net.pragyah.scalgorithms.graphs._

/*
 *  Pearl's Message passing algorithm 
 */

class MessagePassing_SinglyConnected[A](val G:Graph[Variable[A]],P : ProbabilityDistribution[A], val roots:Set[Vertex[Variable[A]]]) {
   
  var aA = Set[Variable[A]]()
  
  assume(G.directed)
  assume(isSinglyConnected)
  
  
  def isSinglyConnected() : boolean = {
    //TODO .... verify if this is singly connected or not
    
    true
  }
  
  def initial_tree(){

	  G.vertices.foreach(X =>{
	    X.tag = VertexTag[A](X.data)
	    X.getEdges.foreach(e =>{
	      val Y = e.getOther(X)
	      e.tag = EdgeTag(X.data,Y.data)
          X.data.values.foreach(x => {
                                  e.tag.asInstanceOf[EdgeTag[A]].pi(x) = 1f
                                })
	               })
	    });

       roots.foreach(root => {
         assume(G.vertices.contains(root))

		       root.data.values.foreach(r => {

		         root.tag.asInstanceOf[VertexTag[A]].P(r) = P(root.data,r,Map()).get // P(r|a) = P(r) for the root
		         root.tag.asInstanceOf[VertexTag[A]].pi(r) = P(root.data,r,Map()).get // pi(r) = P(r) for the root ... since there are no parents and known-children.. there is no one to consult with 
		                                })

		       // for each child X of R ... send pi message
		       root.getEdges.foreach(e => send_pi_msg(root,e.getOther(root)))
       })
       
       
  }
  
  
  def update_tree(V:Variable[A],_v:A) {
    V(Some(_v))
    aA = aA + V
    
    val vVertex = G.vertex(V).get
    
    val vTag = vVertex.tag.asInstanceOf[VertexTag[A]]
    vTag.pi(_v) = 1
    vTag.lambda(_v) = 1
    vTag.P(_v) = 1
    
    V.values.filter(_ != _v).foreach(v => {
			    vTag.pi(v) = 0
			    vTag.lambda(v) = 0
			    vTag.P(v) = 0
   })
    
    
    //for each parent of V that is not in A
    vVertex.inEdges.filter(e => !aA.contains(e.getOther(vVertex).data)).foreach(e => send_lambda_msg(vVertex,e.getOther(vVertex)))

    vVertex.getEdges.foreach(e => {
      val X = e.getOther(vVertex)
        send_pi_msg(vVertex,X)
    })
  }

  //
  // a utility method ....
  //	summation-over-possible-value-combinations-w1,w2,w3 .... wk (where W1...Wk are Y's parents) OF {
  //      P(y|x,w1,w2,....2k) * product-of-all-PI-msgs-by-all-W's-for-given-w1,w2 ....wk (PR-piY(wi)}
  //                                                                  }
  private def parentInfluenceOnYy(X:Vertex[Variable[A]])(x:A)(Y:Vertex[Variable[A]],y:A) : float = {
    
    //return type of getOtherDependencies Map[Set[Tuple2[Variable[A],A]],float] 
    //this one first gets all the other (non-X) dependencies of Y (given value y) given X (and value x) ...

     P.getOtherDependencies(Y.data,y,X.data,x).keySet.toList.foldRight[float](0f)((WwMap,sigma_wi) =>{
       sigma_wi + WwMap.keySet.foldRight[float](1f)((W,piY) => {
                        piY*G.getEdge(G.vertex(W).get,Y).get.tag.asInstanceOf[EdgeTag[A]].pi(WwMap(W))
                     })*P(Y.data,y,WwMap + (X.data -> x)).get 
     })
  }

 
  def send_lambda_msg(Y:Vertex[Variable[A]],X:Vertex[Variable[A]]) {
   
        val edge = X.getEdges.filter(_.getOther(X) == Y).head
        
        //curried function with variable X
        val parentInfluenceOnYyForX = parentInfluenceOnYy(X) _

        X.data.values.foreach(x => {
          // calculate Lambda values for each x ...lambda that has to come up to x from y
// lambda-Y(x)  = summation-over-y(Summation-over-all-W-valsues (P(y|x,w1,w2,...wk)*product-of(all Pi messages for all different combinations of W values)*lambda(y))
            
            // that is ... the opinion that i give to my one parent about his 'that' confiction.. is influenced by 
          // how my convictions (all y) get affected by his 'that' conviction(x) ... along with the convictions of other parents (pi\x).... 
          // each multiplied by the opinion of all my children about my convictions (i.e. lambda value of each y) ... 
            
          
          //curried function .. with value x
          val  parentInfluenceOnYyGiven_Xx = parentInfluenceOnYyForX(x) 
          
          edge.tag.asInstanceOf[EdgeTag[A]].lambda(x) = Y.data.values.foldRight(0f)((y,sigma_y) => {
                                                                                        sigma_y +parentInfluenceOnYyGiven_Xx(Y,y)*Y.tag.asInstanceOf[VertexTag[A]].lambda(y) 
                                                                                    })
          val vtx = X.tag.asInstanceOf[VertexTag[A]]
          vtx.lambda(x) = X.getEdges.foldRight[float](1f)((e,lambda) => (lambda) * e.tag.asInstanceOf[EdgeTag[A]].lambda(x))
          
          vtx.P(x) = vtx.pi(x)*vtx.lambda(x) // this has a hidden alpha .. hence the need to normalize later

        })
        val p = X.tag.asInstanceOf[VertexTag[A]].P 
            //NORMALIZE P(x|a) = P(x1|a) / P(x1|a) +P(x2|a)...+P(xn|a) 
	    val total = p.values./:(0f)(_ + _)
	    p.keySet.foreach(x => {p(x) = p(x)/total})
        
	     //for each parent of X that is not in A
	     X.inEdges.filter(e => !aA.contains(e.getOther(X).data)).foreach(e => {send_lambda_msg(X,e.getOther(X))})     
	     //for each child of X that is not Y
	     X.getEdges.filter(e => e.getOther(X)!= Y && !aA.contains(e.getOther(X).data)).foreach(W => send_pi_msg(X,W.getOther(X)))

  }
  
  private def parentInfluenceOnYy(Y:Vertex[Variable[A]],y:A) : float = {
    
    //return type of getDependencies Map[Set[Tuple2[Variable[A],A]],float] 
    //this one first gets all the  dependencies of Y (given value y) given X (and value x) ...
    // and then calculates the overall influence as
       //summation over all combinations of parent-values ( P(x|parent-values-combinations)*product(piX(each-Parent-Value)) 

     P.getDependencies(Y.data,y).keySet.toList.foldRight[float](0f)((WwMap,sigma_wi) =>{
       sigma_wi + WwMap.keySet.foldRight[float](1f)((W,piY) => {
                        piY*G.getEdge(G.vertex(W).get,Y).get.tag.asInstanceOf[EdgeTag[A]].pi(WwMap(W))
                     })*P(Y.data,y,WwMap).get 
      
     }) 
  }

  
  def send_pi_msg(Z:Vertex[Variable[A]],X:Vertex[Variable[A]]){
    
    //get the edge that connects Z to X (there could be better ways of getting this edge... say directly from the graph ... 
    val edge = G.getEdge(Z,X).get
    
    
    Z.data.values.foreach(z => {//for each possible value z of Z
      // calculate the product of all the Lambda messages coming up from the children of Z (excluding X)
      val lambdaMsgs = Z.getEdges.filter(_.getOther(Z) != X).foldRight[float](1f)((e,lm) => lm*e.tag.asInstanceOf[EdgeTag[A]].lambda(z))

      //calculate the pi message for X for z as a product of pi value of z and the product of lambda messages calculated above
      edge.tag.asInstanceOf[EdgeTag[A]].pi(z) = Z.tag.asInstanceOf[VertexTag[A]].pi(z) * lambdaMsgs

	})


    val vt = X.tag.asInstanceOf[VertexTag[A]]


    
    if(!aA.contains(X.data)){
	    X.data.values.foreach(x => {//for each value x of X
	      
	      //compute X's pi value for x
	         //pi(x) is summation of the prduct of P(x|z) and piX(z) for all z
	//       vt.pi(x)  = Z.data.values.foldRight[float](0f)((z,singlePi) => singlePi + P(X.data,x,Set((Z.data,z)))*edge.tag.asInstanceOf[EdgeTag[A]].pi(z))
	       
	       vt.pi(x)  = parentInfluenceOnYy(X,x)
	       // compute X's probability for x|a
	       vt.P(x) = vt.pi(x)*vt.lambda(x) // this has a hidden alpha .. hence the need to normalize later
	       
	       
	       
		})
	    val p = X.tag.asInstanceOf[VertexTag[A]].P 
	    
	    //NORMALIZE P(x|a) = P(x1|a) / P(x1|a) +P(x2|a)...+P(xn|a) 
	    val total = p.values./:(0f)(_ + _)
	    p.keySet.foreach(x => {p(x) = p(x)/total})
     
	    //for each child Y of X 
	     X.getEdges.foreach(e =>{
	       val Y = e.getOther(X)
	           send_pi_msg(X,Y) // send a pi message to Y
	       }
	     )
    }
        
    // Do not send lambda messages to X's other parents if X and all of X's descendents are uninstantiated.      

    if(!X.data.values.foldRight[boolean](true)((x,lxEq1)=> (lxEq1 && (vt.lambda(x) == 1)))){
      X.inEdges.foreach(e =>{
                           val W = e.getOther(X)
                          if(W != Z && !aA.contains(W.data)){
                            send_lambda_msg(X,W)
                          }
                        })
    }
    
    //notice the above comment ... the reason is taht if X and all X's descendents are unintantiated, X d-separates each of its parents from every other parent.
    //clearly if X and all 's descendents are uninstantiated, then all Xs lambda values are still equal to 1. ... (phrase picked from Neopolitan's book .. chapter 3 page 148)
    //because lambda values are all initialized to 1 in th ebegining and they change only when the lambda message from any child changes .... 
  // which happens only if some descendent is initialized
    
  }
  
  
}



